import numpy as np
import torch

class SystemInitializer:
    def __init__(self, num_clients, total_privacy_budget, base_iterations):
        self.num_clients = num_clients
        self.total_privacy_budget = total_privacy_budget
        self.base_iterations = base_iterations
        
    def server_initialization(self, model_class, *model_args, **model_kwargs):
        """服务器端初始化"""
        global_model = model_class(*model_args, **model_kwargs)
        gradient_buffer = {}
        remaining_budget = self.total_privacy_budget
        clip_threshold = 1.0
        noise_scale = 1.0
        weight_matrix = np.ones(self.num_clients) / self.num_clients
        privacy_window = []
        
        return {
            'global_model': global_model,
            'gradient_buffer': gradient_buffer,
            'remaining_budget': remaining_budget,
            'clip_threshold': clip_threshold,
            'noise_scale': noise_scale,
            'weight_matrix': weight_matrix,
            'privacy_window': privacy_window
        }
    
    def client_initialization(self, client_id, dataset, model_class, *model_args, **model_kwargs):
        """客户端初始化"""
        # 计算本地数据统计量
        if isinstance(dataset.data, torch.Tensor):
            data_tensor = dataset.data
        else:
            data_tensor = torch.tensor(dataset.data)
        mu_k = torch.mean(data_tensor.float(), dim=0)
        
        # 动态设置本地迭代次数
        data_size = len(dataset)
        # 注意：max_data_size需要在外部提供，这里暂时用当前数据集大小
        max_data_size = data_size  # 实际应用中需要从所有客户端数据集中获取最大值
        E_k = max(1, int(self.base_iterations * (data_size / max_data_size)))
        
        # 初始化个性化模型
        personal_model = model_class(*model_args, **model_kwargs)
        
        return {
            'client_id': client_id,
            'local_data': dataset,
            'mu_k': mu_k,
            'E_k': E_k,
            'personal_model': personal_model
    }
